/*
 *  Copyright April 2023
 *  Author: Dong Anakin
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

// Include Files
#include "savitzky_golay.hpp"
#include <Eigen/Core>
#include <iostream>
#define VERIFY_EQUAL(a, b)                                                                         \
    do                                                                                             \
    {                                                                                              \
        auto ret = (a - b).array().abs().maxCoeff();                                               \
        if (ret < 1e-6)                                                                            \
            std::cout << __func__ << ":" << __LINE__ << " OK" << std::endl;                        \
        else                                                                                       \
            std::cerr << __func__ << ":" << __LINE__ << " Failed" << std::endl;                    \
    } while (false)
#define VERIFY_EMPTY(a)                                                                            \
    do                                                                                             \
    {                                                                                              \
        int ret = a.rows() * a.cols();                                                             \
        if (ret == 0)                                                                              \
            std::cout << __func__ << ":" << __LINE__ << " OK" << std::endl;                        \
        else                                                                                       \
            std::cerr << __func__ << ":" << __LINE__ << " Failed" << std::endl;                    \
    } while (false)
void test_correlate01()
{
    Eigen::RowVectorXd input(2);
    input << 1, 2;
    Eigen::RowVectorXd weights(1);
    weights << 2;
    Eigen::RowVectorXd expected(2);
    expected << 2, 4;

    auto r = scipy::ndimage::correlate1d(input, weights);
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, expected);

    auto c = scipy::ndimage::convolve1d(input, weights);
    // savgol_print_matrix(c);

    VERIFY_EQUAL(c, expected);
}
void test_correlate02()
{
    Eigen::RowVectorXd input(3);
    input << 1, 2, 3;
    Eigen::RowVectorXd weights(1);
    weights << 1;
    Eigen::RowVectorXd expected(3);
    expected << 1, 2, 3;

    auto r = scipy::ndimage::correlate1d(input, weights);
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, expected);

    auto c = scipy::ndimage::convolve1d(input, weights);
    // savgol_print_matrix(c);

    VERIFY_EQUAL(c, expected);
}
void test_correlate03()
{
    Eigen::RowVectorXd input(1);
    input << 1;
    Eigen::RowVectorXd weights(2);
    weights << 1, 1;
    Eigen::RowVectorXd expected(1);
    expected << 2;

    auto r = scipy::ndimage::correlate1d(input, weights);
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, expected);

    auto c = scipy::ndimage::convolve1d(input, weights);
    // savgol_print_matrix(c);

    VERIFY_EQUAL(c, expected);
}

void test_correlate04()
{
    Eigen::RowVectorXd input(2);
    input << 1, 2;
    Eigen::RowVectorXd weights(2);
    weights << 1, 1;
    Eigen::RowVectorXd tcor(2);
    tcor << 2, 3;
    Eigen::RowVectorXd tcov(2);
    tcov << 3, 4;

    auto r = scipy::ndimage::correlate1d(input, weights);
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, tcor);

    auto c = scipy::ndimage::convolve1d(input, weights);
    // savgol_print_matrix(c);

    VERIFY_EQUAL(c, tcov);
}

void test_correlate05()
{
    Eigen::RowVectorXd input(3);
    input << 1, 2, 3;
    Eigen::RowVectorXd weights(2);
    weights << 1, 1;
    Eigen::RowVectorXd tcor(3);
    tcor << 2, 3, 5;
    Eigen::RowVectorXd tcov(3);
    tcov << 3, 5, 6;

    auto r = scipy::ndimage::correlate1d(input, weights);
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, tcor);

    auto c = scipy::ndimage::convolve1d(input, weights);
    // savgol_print_matrix(c);

    VERIFY_EQUAL(c, tcov);
}

void test_correlate06()
{
    Eigen::RowVectorXd input(3);
    input << 1, 2, 3;
    Eigen::RowVectorXd weights(3);
    weights << 1, 2, 3;
    Eigen::RowVectorXd tcor(3);
    tcor << 9, 14, 17;
    Eigen::RowVectorXd tcov(3);
    tcov << 7, 10, 15;

    auto r = scipy::ndimage::correlate1d(input, weights);
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, tcor);

    auto c = scipy::ndimage::convolve1d(input, weights);
    // savgol_print_matrix(c);

    VERIFY_EQUAL(c, tcov);
}

void test_correlate07()
{
    Eigen::RowVectorXd input(3);
    input << 1, 2, 3;
    Eigen::RowVectorXd weights(3);
    weights << 1, 2, 1;
    Eigen::RowVectorXd expected(3);
    expected << 5, 8, 11;

    auto r = scipy::ndimage::correlate1d(input, weights);
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, expected);

    auto c = scipy::ndimage::convolve1d(input, weights);
    // savgol_print_matrix(c);

    VERIFY_EQUAL(c, expected);
}

void test_correlate08()
{
    Eigen::RowVectorXd input(3);
    input << 1, 2, 3;
    Eigen::RowVectorXd weights(3);
    weights << 1, 2, -1;
    Eigen::RowVectorXd tcor(3);
    tcor << 1, 2, 5;
    Eigen::RowVectorXd tcov(3);
    tcov << 3, 6, 7;

    auto r = scipy::ndimage::correlate1d(input, weights);
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, tcor);

    auto c = scipy::ndimage::convolve1d(input, weights);
    // savgol_print_matrix(c);

    VERIFY_EQUAL(c, tcov);
}

void test_correlate09()
{
    Eigen::RowVectorXd input(0);
    Eigen::RowVectorXd weights(2);
    weights << 1, 1;
    Eigen::RowVectorXd tcor(0);
    Eigen::RowVectorXd tcov(0);

    auto r = scipy::ndimage::correlate1d(input, weights);
    // savgol_print_matrix(r);
    VERIFY_EMPTY(r);

    auto c = scipy::ndimage::convolve1d(input, weights);
    // savgol_print_matrix(c);

    VERIFY_EMPTY(c);
}

void test_correlate17()
{
    Eigen::RowVectorXd input(3);
    input << 1, 2, 3;
    Eigen::RowVectorXd weights(2);
    weights << 1, 1;
    Eigen::RowVectorXd tcor(3);
    tcor << 3, 5, 6;
    Eigen::RowVectorXd tcov(3);
    tcov << 2, 3, 5;

    auto r = scipy::ndimage::correlate1d(input, weights, "reflect", 0.0f, -1);
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, tcor);

    auto c = scipy::ndimage::convolve1d(input, weights, "reflect", 0.0f, -1);
    // savgol_print_matrix(c);

    VERIFY_EQUAL(c, tcov);
}

void test_correlate22()
{
    Eigen::RowVectorXd input(6);
    input << 1, 2, 3, 2, 4, 6;
    Eigen::RowVectorXd weights(3);
    weights << 1, 2, 1;
    Eigen::RowVectorXd expected(6);
    expected << 10, 8, 10, 11, 16, 17;

    auto r = scipy::ndimage::correlate1d(input, weights, "wrap");
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, expected);

    auto c = scipy::ndimage::convolve1d(input, weights, "wrap");
    // savgol_print_matrix(c);
    VERIFY_EQUAL(c, expected);
}

void test_correlate23()
{
    Eigen::RowVectorXd input(6);
    input << 1, 2, 3, 2, 4, 6;
    Eigen::RowVectorXd weights(3);
    weights << 1, 2, 1;
    Eigen::RowVectorXd expected(6);
    expected << 5, 8, 10, 11, 16, 22;

    auto r = scipy::ndimage::correlate1d(input, weights, "nearest");
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, expected);

    auto c = scipy::ndimage::convolve1d(input, weights, "nearest");
    // savgol_print_matrix(c);

    VERIFY_EQUAL(c, expected);
}

void test_correlate24()
{
    Eigen::RowVectorXd input(6);
    input << 1, 2, 3, 2, 4, 6;
    Eigen::RowVectorXd weights(3);
    weights << 1, 2, 1;
    Eigen::RowVectorXd tcor(6);
    tcor << 8, 10, 11, 16, 22, 24;
    Eigen::RowVectorXd tcov(6);
    tcov << 4, 5, 8, 10, 11, 16;

    auto r = scipy::ndimage::correlate1d(input, weights, "nearest", 0.0f, -1);
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, tcor);

    auto c = scipy::ndimage::convolve1d(input, weights, "nearest", 0.0f, -1);
    // savgol_print_matrix(c);

    VERIFY_EQUAL(c, tcov);
}

void test_correlate25()
{
    Eigen::RowVectorXd input(6);
    input << 1, 2, 3, 2, 4, 6;
    Eigen::RowVectorXd weights(3);
    weights << 1, 2, 1;
    Eigen::RowVectorXd tcor(6);
    tcor << 4, 5, 8, 10, 11, 16;
    Eigen::RowVectorXd tcov(6);
    tcov << 8, 10, 11, 16, 22, 24;

    auto r = scipy::ndimage::correlate1d(input, weights, "nearest", 0.0f, 1);
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, tcor);

    auto c = scipy::ndimage::convolve1d(input, weights, "nearest", 0.0f, 1);
    // savgol_print_matrix(c);

    VERIFY_EQUAL(c, tcov);
}

void test_correlate26()
{
    Eigen::RowVectorXd input(1);
    input << 1;
    Eigen::RowVectorXd weights(5);
    weights << 1, 1, 1, 1, 1;
    Eigen::RowVectorXd expected(1);
    expected << 5;

    auto r = scipy::ndimage::correlate1d(input, weights, "mirror");
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, expected);

    auto c = scipy::ndimage::convolve1d(input, weights, "mirror");
    // savgol_print_matrix(c);

    VERIFY_EQUAL(c, expected);
}

void test_correlate27()
{
    Eigen::RowVectorXd input(6);
    input << 1, 2, 3, 2, 4, 6;
    Eigen::RowVectorXd weights(3);
    weights << 1, 2, 1;
    Eigen::RowVectorXd tcor(6);
    tcor << 9, 8, 10, 11, 16, 21;
    Eigen::RowVectorXd tcov(6);
    tcov << 9, 8, 10, 11, 16, 21;

    auto r = scipy::ndimage::correlate1d(input, weights, "constant", 5.0f);
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, tcor);

    auto c = scipy::ndimage::convolve1d(input, weights, "constant", 5.0f);
    // savgol_print_matrix(c);

    VERIFY_EQUAL(c, tcov);
}

void test_correlate28()
{
    Eigen::RowVectorXd input(6);
    input << 1, 2, 3, 2, 4, 6;
    Eigen::RowVectorXd weights(3);
    weights << 1, 2, -1;
    Eigen::RowVectorXd tcor(6);
    tcor << 5, 2, 6, 3, 4, 11;
    Eigen::RowVectorXd tcov(6);
    tcov << -1, 6, 6, 5, 12, 13;

    auto r = scipy::ndimage::correlate1d(input, weights, "constant", 5.0f);
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, tcor);

    auto c = scipy::ndimage::convolve1d(input, weights, "constant", 5.0f);
    // savgol_print_matrix(c);
    VERIFY_EQUAL(c, tcov);
}

void test_correlate30()
{
    Eigen::RowVectorXd input(3);
    input << 1, 2, 3;
    Eigen::RowVectorXd weights(8);
    weights << 1, 0, 0, 0, 0, 0, 0, 0;
    Eigen::RowVectorXd tcor(3);
    tcor << 1, 1, 1;
    Eigen::RowVectorXd tcov(3);
    tcov << 3, 3, 3;

    auto r = scipy::ndimage::correlate1d(input, weights, "nearest");
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, tcor);

    auto c = scipy::ndimage::convolve1d(input, weights, "nearest");
    // savgol_print_matrix(c);
    VERIFY_EQUAL(c, tcov);
}

void test_correlate31()
{
    Eigen::RowVectorXd input(3);
    input << 1, 2, 3;
    Eigen::RowVectorXd weights(8);
    weights << 1, 0, 0, 0, 0, 0, 0, 0;
    Eigen::RowVectorXd tcor(3);
    tcor << 3, 1, 2;
    Eigen::RowVectorXd tcov(3);
    tcov << 2, 3, 1;

    auto r = scipy::ndimage::correlate1d(input, weights, "wrap");
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, tcor);

    auto c = scipy::ndimage::convolve1d(input, weights, "wrap");
    // savgol_print_matrix(c);
    VERIFY_EQUAL(c, tcov);
}

void test_correlate32()
{
    Eigen::RowVectorXd input(3);
    input << 1, 2, 3;
    Eigen::RowVectorXd weights(8);
    weights << 1, 0, 0, 0, 0, 0, 0, 0;
    Eigen::RowVectorXd tcor(3);
    tcor << 3, 3, 2;
    Eigen::RowVectorXd tcov(3);
    tcov << 2, 1, 1;

    auto r = scipy::ndimage::correlate1d(input, weights, "reflect");
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, tcor);

    auto c = scipy::ndimage::convolve1d(input, weights, "reflect");
    // savgol_print_matrix(c);
    VERIFY_EQUAL(c, tcov);
}

void test_correlate33()
{
    Eigen::RowVectorXd input(3);
    input << 1, 2, 3;
    Eigen::RowVectorXd weights(8);
    weights << 1, 0, 0, 0, 0, 0, 0, 0;
    Eigen::RowVectorXd tcor(3);
    tcor << 1, 2, 3;
    Eigen::RowVectorXd tcov(3);
    tcov << 1, 2, 3;

    auto r = scipy::ndimage::correlate1d(input, weights, "mirror");
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, tcor);

    auto c = scipy::ndimage::convolve1d(input, weights, "mirror");
    // savgol_print_matrix(c);
    VERIFY_EQUAL(c, tcov);
}

void test_correlate34()
{
    Eigen::RowVectorXd input(3);
    input << 1, 2, 3;
    Eigen::RowVectorXd weights(8);
    weights << 1, 0, 0, 0, 0, 0, 0, 0;
    Eigen::RowVectorXd tcor(3);
    tcor << 0, 0, 0;
    Eigen::RowVectorXd tcov(3);
    tcov << 0, 0, 0;

    auto r = scipy::ndimage::correlate1d(input, weights, "constant");
    // savgol_print_matrix(r);
    VERIFY_EQUAL(r, tcor);

    auto c = scipy::ndimage::convolve1d(input, weights, "constant");
    // savgol_print_matrix(c);
    VERIFY_EQUAL(c, tcov);
}

void test_savgol_coeffs_01()
{
    Eigen::RowVectorXd expected(5);
    expected << 0.77142857, -0.18571429, -0.57142857, -0.38571429, 0.37142857;

    int pos    = 4;
    auto coeff = scipy::signal::savgol_coeffs(5, 2, 1, 1.0f, &pos, "conv");
    // savgol_print_matrix(coeff);
    VERIFY_EQUAL(coeff, expected);
}

void test_savgol_coeffs_02()
{
    Eigen::RowVectorXd expected(5);
    expected << 0.37142857, -0.38571429, -0.57142857, -0.18571429, 0.77142857;

    int pos    = 4;
    auto coeff = scipy::signal::savgol_coeffs(5, 2, 1, 1.0f, &pos, "dot");
    // savgol_print_matrix(coeff);
    VERIFY_EQUAL(coeff, expected);
}

void test_savgol_filter_01()
{
    Eigen::RowVectorXd x(9);
    x << 2, 2, 5, 2, 1, 0, 1, 4, 9;

    Eigen::RowVectorXd expected(9);
    expected << 1.65714286, 3.17142857, 3.54285714, 2.85714286, 0.65714286, 0.17142857, 1.0, 4.0,
        9.0;

    auto y = scipy::signal::savgol_filter(x, 5, 2);
    // savgol_print_matrix(y);
    VERIFY_EQUAL(y, expected);
}

void test_savgol_filter_02()
{
    Eigen::RowVectorXd x(9);
    x << 2, 2, 5, 2, 1, 0, 1, 4, 9;

    Eigen::RowVectorXd expected(9);
    expected << 1.48571429, 3.02857143, 3.54285714, 2.85714286, 0.65714286, 0.17142857, 1.,
        5.02857143, 6.94285714;

    auto y = scipy::signal::savgol_filter(x, 5, 2, 0, 1.0, "mirror");
    // savgol_print_matrix(y);
    VERIFY_EQUAL(y, expected);
}

void test_savgol_filter_03()
{
    Eigen::RowVectorXd x(9);
    x << 2, 2, 5, 2, 1, 0, 1, 4, 9;

    Eigen::RowVectorXd expected(9);
    expected << 1.71428571, -0.42857143, -1.14285714, -0.85714286, 1.14285714, 1.42857143, 2.,
        -1.42857143, -3.14285714;

    auto y = scipy::signal::savgol_filter(x, 5, 3, 2, 1.0, "mirror");
    // savgol_print_matrix(y);
    VERIFY_EQUAL(y, expected);
}

void test_savgol_filter_04()
{
    Eigen::RowVectorXd x(9);
    x << 2, 2, 5, 2, 1, 0, 1, 4, 9;

    Eigen::RowVectorXd expected(9);
    expected << -0.14285714, -0.64285714, -1.14285714, -0.85714286, 1.14285714, 1.42857143, 2., 2.,
        2.;

    auto y = scipy::signal::savgol_filter(x, 5, 3, 2, 1.0);
    // savgol_print_matrix(y);
    VERIFY_EQUAL(y, expected);
}

void test_numpy_polyfit_01()
{
    Eigen::RowVectorXd x(6);
    x << 0.0, 1.0, 2.0, 3.0, 4.0, 5.0;
    Eigen::RowVectorXd y(6);
    y << 0.0, 0.8, 0.9, 0.1, -0.8, -1.0;
    Eigen::RowVectorXd expected(4);
    expected << 0.08703704, -0.81349206, 1.69312169, -0.03968254;

    Eigen::RowVectorXd z = numpy::polyfit(x, y, 3);
    // savgol_print_matrix(z);
    VERIFY_EQUAL(z, expected);
}

int main(int argc, char **argv)
{
    test_correlate01();
    test_correlate02();
    test_correlate03();
    test_correlate04();
    test_correlate05();
    test_correlate06();
    test_correlate07();
    test_correlate08();
    test_correlate09();
    test_correlate17();
    test_correlate22();
    test_correlate23();
    test_correlate24();
    test_correlate25();
    test_correlate26();
    test_correlate27();
    test_correlate28();
    test_correlate30();
    test_correlate31();
    test_correlate32();
    test_correlate33();
    test_correlate34();
    test_numpy_polyfit_01();
    test_savgol_coeffs_01();
    test_savgol_coeffs_02();
    test_savgol_filter_01();
    test_savgol_filter_02();
    test_savgol_filter_03();
    test_savgol_filter_04();
    return 0;
}
